#coding=utf8
import torch
from torch import nn
from torch.autograd import Variable
import math
import numpy as np
from models.other_layers import grouplinear_clf
from torch.utils.data import TensorDataset, DataLoader
from torch.optim import lr_scheduler
import torch.optim as optim
from torch.nn import CrossEntropyLoss
import torch.nn.functional as F
from sklearn.metrics import f1_score, precision_recall_fscore_support, classification_report
import logging
from utils.train import dt
from math import ceil
import copy
from utils.ensemble_tools import class_to_idx_map, idx_to_class_map
from utils.ensemble_tools import simple_npy_predict, class_to_idx_map, simple_train_npy
import os
import pandas as pd


true_labels = np.load('/home/gserver/zhangchi/tianwen/val/val-labels.npy')

all_score = np.zeros((48386, 5, 4),dtype=np.float32)
val_score1 = np.load('/home/gserver/zhangchi/tianwen/val/resnet20_crop_aug8-0.8229.npy')
val_score2 = np.load('/home/gserver/zhangchi/tianwen/val/xception-0.8138.npy')
val_score3 = np.load('/home/gserver/zhangchi/tianwen/val/xception_crop-0.7995-aug8-0.8129.npy')
val_score4 = np.load('/home/gserver/zhangchi/tianwen/val/resnet20_crop_smpunknown-0.8226.npy')
val_score5 = np.load('/home/gserver/zhangchi/tianwen/val/resnet20_smpunknown1-0.8083-aug8-0.8133.npy')



all_score[:,0,:] = val_score1
all_score[:,1,:] = val_score2
all_score[:,2,:] = val_score3
all_score[:,3,:] = val_score4
all_score[:,4,:] = val_score5

x_train = all_score
y_train = true_labels
x_val = all_score
y_val = true_labels


# model prepare
resume = None
model = grouplinear_clf(group=5, classes=4, bias=False)
# model = nn.Linear(in_features=12, out_features=4, bias=True)
model.weight.data = torch.from_numpy(np.zeros((5,1))+0.2).float()
print model.weight.data
model = torch.nn.DataParallel(model)
if resume:
    model.load_state_dict(torch.load(resume))

model = model.cuda()


optimizer = optim.SGD(model.parameters(), lr=0.0011, momentum=0.9, weight_decay=0)
criterion = CrossEntropyLoss()
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.5)


best_model_wts, best_acc = simple_train_npy(model, x_train, y_train, x_val, y_val,
                                            start_epoch=0, epoch_num=20,
                                            optimizer=optimizer,
                                            criterion=criterion,
                                            exp_lr_scheduler=exp_lr_scheduler,
                                            bs_train=32*3, bs_val=32*3, val_inter=200)

print 'best f1:',best_acc
save_path = './merge_weights/linear-weights-%.5f'%best_acc
torch.save(best_model_wts, save_path)

model.load_state_dict(best_model_wts)
#
#
#
#
# # pred test
#
# all_score = np.zeros((100000, 3, 4),dtype=np.float32)
# val_score1 = np.load('/home/gserver/zhangchi/tianwen/online_pred/scores/resnet20_crop-0.8138-aug8-0.8229.npy')
# val_score2 = np.load('/home/gserver/zhangchi/tianwen/online_pred/scores/xception-0.8138.npy')
# val_score3 = np.load('/home/gserver/zhangchi/tianwen/online_pred/scores/xception_crop-0.7995-aug8-0.8129.npy')
#
# all_score[:,0,:] = val_score1
# all_score[:,1,:] = val_score2
# all_score[:,2,:] = val_score3
#
# labels = np.zeros(all_score.shape[0],dtype=int)
#
#
# x_test = all_score
# _, online_preds, _ = simple_npy_predict(model, x_test, labels, 32 * 3, usecuda=1)
#
#
# # save csv file to submit
#
# rawdata_root = '/media/gserver/data/tianwen/rawdata'
# test_index = pd.read_csv(os.path.join(rawdata_root, 'first_test_index_20180131.csv'))
# sub = pd.DataFrame({'id':test_index['id'].tolist(),
#                     'type':online_preds})
# sub['type'] = sub['type'].apply(lambda x: idx_to_class_map[x])
#
# # sub[['id','type']].to_csv('./merge/merge_sub/gl-weights-%.5f.csv'%(best_acc),index=False,header=None)
# print sub
# print sub['type'].value_counts()
